package com.atguigu.ai.chat.controller;

import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.chat.prompt.PromptTemplate;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.*;
import reactor.core.publisher.Flux;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;

@RestController
@RequestMapping("/prompt")
public class PromptController {
    @Autowired
    private ChatClient chatClient;

    /**
     * 基础提示词操作 -只有用户消息
     *
     * @param message
     * @return
     */
    @GetMapping("/basic")
    public String basicPrompt(@RequestParam(value = "msg", defaultValue = "你好") String message) {
        return chatClient.prompt().user(message).call().content();
    }

    /**
     * 系统提示词+用户提示词
     */
    @GetMapping("/system")
    public String systemPrompt(@RequestParam(value = "system", defaultValue = "你是一个专业的Java开发架构师") String system,
                               @RequestParam(value = "user", defaultValue = "请介绍下Spring AI") String userMessage) {
        return chatClient.prompt().system(system).user(userMessage).call().content();
    }

    @GetMapping("/template")
    public String templatePrompt(@RequestParam(value = "topic", defaultValue = "Spring AI") String topic,
                                 @RequestParam(value = "level", defaultValue = "初级") String level) {
        //定义提示词模板
        String templateText = """
                你是一名经验丰富的技术讲师。
                                请为{level}学员详细介绍{topic}的相关知识。
                                要求：
                                1. 语言通俗易懂
                                2. 包含实际应用场景
                                3. 提供学习建议
                """;
        PromptTemplate promptTemplate = new PromptTemplate(templateText);
        Prompt prompt = promptTemplate.create(Map.of("level", level, "topic", topic));
        return chatClient.prompt(prompt).call().content();
    }

    /**
     * 多轮对话提示词
     */
    @PostMapping("/conversation")
    public String conversationPrompt(@RequestBody ConversationRequest request) {
        List<Message> messages = new ArrayList<>();
        // 添加系统消息
        if (request.getSystemMessage() != null && !request.getSystemMessage().isEmpty()) {
            messages.add(new SystemMessage(request.getSystemMessage()));
        }
        //添加历史对话
        if (request.getHistory() != null) {
            for (ConversationRequest.ChatMessage chatMessage : request.getHistory()) {
                if ("user".equals(chatMessage.getRole())) {
                    messages.add(new UserMessage(chatMessage.getContent()));
                } else if ("system".equals(chatMessage.getRole())) {
                    messages.add(new SystemMessage(chatMessage.getContent()));
                }
            }
        }
        messages.add(new UserMessage(request.getCurrentMessage()));
        Prompt prompt = new Prompt(messages);
        return chatClient.prompt(prompt).call().content();
    }

    /**
     * 流式提示词响应
     */
    @GetMapping(value = "/stream", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<String> streamPrompt(
            @RequestParam(value = "system", defaultValue = "你是一个友好的AI助手") String systemMessage,
            @RequestParam(value = "user", defaultValue = "请介绍人工智能的发展历史") String userMessage) {

        return chatClient.prompt()
                .system(systemMessage)
                .user(userMessage)
                .stream()
                .content();
    }

    // 内部类定义请求对象
    public static class ConversationRequest {
        private String systemMessage;
        private List<ChatMessage> history;
        private String currentMessage;

        // getter和setter方法
        public String getSystemMessage() {
            return systemMessage;
        }

        public void setSystemMessage(String systemMessage) {
            this.systemMessage = systemMessage;
        }

        public List<ChatMessage> getHistory() {
            return history;
        }

        public void setHistory(List<ChatMessage> history) {
            this.history = history;
        }

        public String getCurrentMessage() {
            return currentMessage;
        }

        public void setCurrentMessage(String currentMessage) {
            this.currentMessage = currentMessage;
        }

        public static class ChatMessage {
            private String role; // "user" 或 "assistant"
            private String content;

            public String getRole() {
                return role;
            }

            public void setRole(String role) {
                this.role = role;
            }

            public String getContent() {
                return content;
            }

            public void setContent(String content) {
                this.content = content;
            }
        }
    }
}
